package com.example.sse.test.util;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.apache.http.HttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * SSE 测试客户端工具类
 * 提供方便的 SSE 连接测试功能
 */
public class SseTestClient {
    
    private final String baseUrl;
    private CloseableHttpClient httpClient;
    private boolean isConnected = false;
    private CompletableFuture<Void> readTask;
    
    public SseTestClient(String baseUrl) {
        this.baseUrl = baseUrl;
        this.httpClient = HttpClients.createDefault();
    }
    
    /**
     * 连接到 SSE 端点
     * @param clientId 客户端ID（可选）
     * @return 连接是否成功
     */
    public boolean connect(String clientId) {
        try {
            String url = baseUrl + "/api/sse/connect";
            if (clientId != null && !clientId.trim().isEmpty()) {
                url += "?clientId=" + clientId;
            }
            
            HttpGet request = new HttpGet(url);
            request.setHeader("Accept", "text/event-stream");
            request.setHeader("Cache-Control", "no-cache");
            
            HttpResponse response = httpClient.execute(request);
            
            if (response.getStatusLine().getStatusCode() == 200) {
                isConnected = true;
                System.out.println("SSE 连接建立成功: " + url);
                return true;
            } else {
                System.err.println("SSE 连接失败，状态码: " + response.getStatusLine().getStatusCode());
                return false;
            }
        } catch (Exception e) {
            System.err.println("连接 SSE 时发生错误: " + e.getMessage());
            return false;
        }
    }
    
    /**
     * 监听 SSE 消息
     * @param messageHandler 消息处理器
     * @param maxMessages 最大消息数量（0表示无限制）
     * @param timeoutSeconds 超时时间（秒）
     */
    public void listenMessages(Consumer<SseMessage> messageHandler, int maxMessages, int timeoutSeconds) {
        if (!isConnected) {
            throw new IllegalStateException("未建立连接，请先调用 connect() 方法");
        }
        
        readTask = CompletableFuture.runAsync(() -> {
            try {
                String url = baseUrl + "/api/sse/connect";
                HttpGet request = new HttpGet(url);
                request.setHeader("Accept", "text/event-stream");
                
                HttpResponse response = httpClient.execute(request);
                
                if (response.getStatusLine().getStatusCode() == 200) {
                    try (BufferedReader reader = new BufferedReader(
                            new InputStreamReader(response.getEntity().getContent()))) {
                        
                        String line;
                        int messageCount = 0;
                        SseMessage currentMessage = new SseMessage();
                        
                        while ((line = reader.readLine()) != null) {
                            if (maxMessages > 0 && messageCount >= maxMessages) {
                                break;
                            }
                            
                            if (line.trim().isEmpty()) {
                                // 空行表示消息结束
                                if (currentMessage.hasData()) {
                                    messageHandler.accept(currentMessage);
                                    messageCount++;
                                    currentMessage = new SseMessage();
                                }
                            } else if (line.startsWith("data:")) {
                                currentMessage.setData(line.substring(5).trim());
                            } else if (line.startsWith("event:")) {
                                currentMessage.setEvent(line.substring(6).trim());
                            } else if (line.startsWith("id:")) {
                                currentMessage.setId(line.substring(3).trim());
                            } else if (line.startsWith("retry:")) {
                                currentMessage.setRetry(line.substring(6).trim());
                            }
                        }
                    }
                }
            } catch (IOException e) {
                System.err.println("读取 SSE 消息时发生错误: " + e.getMessage());
            }
        });
        
        // 设置超时
        if (timeoutSeconds > 0) {
            try {
                readTask.get(timeoutSeconds, TimeUnit.SECONDS);
            } catch (Exception e) {
                System.out.println("监听超时或完成");
            }
        }
    }
    
    /**
     * 获取连接状态
     */
    public JSONObject getConnectionStatus() {
        try {
            HttpGet request = new HttpGet(baseUrl + "/api/sse/status");
            request.setHeader("Accept", "application/json");
            
            HttpResponse response = httpClient.execute(request);
            
            if (response.getStatusLine().getStatusCode() == 200) {
                try (BufferedReader reader = new BufferedReader(
                        new InputStreamReader(response.getEntity().getContent()))) {
                    
                    StringBuilder responseBody = new StringBuilder();
                    String line;
                    while ((line = reader.readLine()) != null) {
                        responseBody.append(line);
                    }
                    
                    return JSON.parseObject(responseBody.toString());
                }
            }
        } catch (Exception e) {
            System.err.println("获取状态时发生错误: " + e.getMessage());
        }
        return null;
    }
    
    /**
     * 收集指定数量的消息
     * @param count 消息数量
     * @param timeoutSeconds 超时时间
     * @return 消息列表
     */
    public List<SseMessage> collectMessages(int count, int timeoutSeconds) {
        List<SseMessage> messages = new ArrayList<>();
        
        listenMessages(messages::add, count, timeoutSeconds);
        
        return messages;
    }
    
    /**
     * 关闭连接
     */
    public void close() {
        try {
            if (readTask != null && !readTask.isDone()) {
                readTask.cancel(true);
            }
            if (httpClient != null) {
                httpClient.close();
            }
            isConnected = false;
        } catch (IOException e) {
            System.err.println("关闭连接时发生错误: " + e.getMessage());
        }
    }
    
    /**
     * SSE 消息模型
     */
    public static class SseMessage {
        private String data;
        private String event;
        private String id;
        private String retry;
        private long receivedTime;
        
        public SseMessage() {
            this.receivedTime = System.currentTimeMillis();
        }
        
        public boolean hasData() {
            return data != null && !data.trim().isEmpty();
        }
        
        public JSONObject getDataAsJson() {
            if (data != null) {
                try {
                    return JSON.parseObject(data);
                } catch (Exception e) {
                    return null;
                }
            }
            return null;
        }
        
        // Getters and Setters
        public String getData() { return data; }
        public void setData(String data) { this.data = data; }
        public String getEvent() { return event; }
        public void setEvent(String event) { this.event = event; }
        public String getId() { return id; }
        public void setId(String id) { this.id = id; }
        public String getRetry() { return retry; }
        public void setRetry(String retry) { this.retry = retry; }
        public long getReceivedTime() { return receivedTime; }
        
        @Override
        public String toString() {
            return String.format("SseMessage{data='%s', event='%s', id='%s', time=%d}", 
                    data, event, id, receivedTime);
        }
    }
}